{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gJ1C0j6bM8Ns"
      },
      "source": [
        "Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SCcCMGl4NOE0"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wVhbbTC_NYJq"
      },
      "source": [
        "# Recommending movies: retrieval using a sequential model\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/recommenders/examples/sequential_retrieval\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/recommenders/blob/main/docs/examples/sequential_retrieval.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/recommenders/blob/main/docs/examples/sequential_retrieval.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/recommenders/docs/examples/sequential_retrieval.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FSQBHTghPGcg"
      },
      "source": [
        "In this tutorial, we are going to build a sequential retrieval model. Sequential recommendation is a popular model that looks at a sequence of  items that users have interacted with previously and then predicts the next item. Here the order of the items within each sequence matters, so we are going to use a recurrent neural network to model the sequential relationship. For more details, please refer to this [GRU4Rec paper](https://arxiv.org/abs/1511.06939).\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "447R2K8kRY0Z"
      },
      "source": [
        "## Imports\n",
        "\n",
        "First let's get our dependencies and imports out of the way."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NM3uchF8deZL"
      },
      "outputs": [],
      "source": [
        "!pip install -q tensorflow-recommenders\n",
        "!pip install -q --upgrade tensorflow-datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HfpLMl1tc4OE"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import pprint\n",
        "import tempfile\n",
        "\n",
        "from typing import Dict, Text\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow_recommenders as tfrs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GULCYkofR2pP"
      },
      "source": [
        "## Preparing the dataset\n",
        "\n",
        "Next, we need to prepare our dataset. We are going to leverage the [data generation utility](https://github.com/tensorflow/examples/blob/master/lite/examples/recommendation/ml/data/example_generation_movielens.py) in this [TensorFlow Lite On-device Recommendation reference app](https://www.tensorflow.org/lite/examples/recommendation/overview).\n",
        "\n",
        "MovieLens 1M data contains ratings.dat (*columns: UserID, MovieID, Rating, Timestamp*), and movies.dat (*columns: MovieID, Title, Genres*). The example generation script download the 1M dataset, takes both files, only keep ratings higher than 2, form user movie interaction timelines, sample activities as labels and 10 previous user activities as the context for prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sJiKu6hRc2H2"
      },
      "outputs": [],
      "source": [
        "!wget -nc https://raw.githubusercontent.com/tensorflow/examples/master/lite/examples/recommendation/ml/data/example_generation_movielens.py\n",
        "!python -m example_generation_movielens  --data_dir=data/raw  --output_dir=data/examples  --min_timeline_length=3  --max_context_length=10  --max_context_movie_genre_length=10  --min_rating=2  --train_data_fraction=0.9  --build_vocabs=False"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U5WEdiUNSmNb"
      },
      "source": [
        "Here is a sample of the generated dataset.\n",
        "\n",
        "```\n",
        "0 : {\n",
        "  features: {\n",
        "    feature: {\n",
        "      key  : \"context_movie_id\"\n",
        "      value: { int64_list: { value: [ 1124, 2240, 3251, ..., 1268 ] } }\n",
        "    }\n",
        "    feature: {\n",
        "      key  : \"context_movie_rating\"\n",
        "      value: { float_list: {value: [ 3.0, 3.0, 4.0, ..., 3.0 ] } }\n",
        "    }\n",
        "    feature: {\n",
        "      key  : \"context_movie_year\"\n",
        "      value: { int64_list: { value: [ 1981, 1980, 1985, ..., 1990 ] } }\n",
        "    }\n",
        "    feature: {\n",
        "      key  : \"context_movie_genre\"\n",
        "      value: { bytes_list: { value: [ \"Drama\", \"Drama\", \"Mystery\", ..., \"UNK\" ] } }\n",
        "    }\n",
        "    feature: {\n",
        "      key  : \"label_movie_id\"\n",
        "      value: { int64_list: { value: [ 3252 ] }  }\n",
        "    }\n",
        "  }\n",
        "}\n",
        "```\n",
        "You can see that it includes a sequence of context movie IDs, and a label movie ID (next movie), plus context features such as movie year, rating and genre. \n",
        "\n",
        "In our case we will only be using the sequence of context movie IDs and the label movie ID. You can refer to the [Leveraging context features tutorial](https://www.tensorflow.org/recommenders/examples/context_features) to learn more about adding additional context features."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JQ2aLaCndSYP"
      },
      "outputs": [],
      "source": [
        "train_filename = \"./data/examples/train_movielens_1m.tfrecord\"\n",
        "train = tf.data.TFRecordDataset(train_filename)\n",
        "\n",
        "test_filename = \"./data/examples/test_movielens_1m.tfrecord\"\n",
        "test = tf.data.TFRecordDataset(test_filename)\n",
        "\n",
        "feature_description = {\n",
        "    'context_movie_id': tf.io.FixedLenFeature([10], tf.int64, default_value=np.repeat(0, 10)),\n",
        "    'context_movie_rating': tf.io.FixedLenFeature([10], tf.float32, default_value=np.repeat(0, 10)),\n",
        "    'context_movie_year': tf.io.FixedLenFeature([10], tf.int64, default_value=np.repeat(1980, 10)),\n",
        "    'context_movie_genre': tf.io.FixedLenFeature([10], tf.string, default_value=np.repeat(\"Drama\", 10)),\n",
        "    'label_movie_id': tf.io.FixedLenFeature([1], tf.int64, default_value=0),\n",
        "}\n",
        "\n",
        "def _parse_function(example_proto):\n",
        "  return tf.io.parse_single_example(example_proto, feature_description)\n",
        "\n",
        "train_ds = train.map(_parse_function).map(lambda x: {\n",
        "    \"context_movie_id\": tf.strings.as_string(x[\"context_movie_id\"]),\n",
        "    \"label_movie_id\": tf.strings.as_string(x[\"label_movie_id\"])\n",
        "})\n",
        "\n",
        "test_ds = test.map(_parse_function).map(lambda x: {\n",
        "    \"context_movie_id\": tf.strings.as_string(x[\"context_movie_id\"]),\n",
        "    \"label_movie_id\": tf.strings.as_string(x[\"label_movie_id\"])\n",
        "})\n",
        "\n",
        "for x in train_ds.take(1).as_numpy_iterator():\n",
        "  pprint.pprint(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OKmUecCVTf9Z"
      },
      "source": [
        "Now our train/test datasets include only a sequence of historical movie IDs and a label of next movie ID. Note that we use `[10]` as the shape of the features during tf.Example parsing because we specify 10 as the length of context features in the example generateion step.\n",
        "\n",
        "We need one more thing before we can start building the model - the vocabulary for our movie IDs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Szgmv8fvc7Tj"
      },
      "outputs": [],
      "source": [
        "movies = tfds.load(\"movielens/1m-movies\", split='train')\n",
        "movies = movies.map(lambda x: x[\"movie_id\"])\n",
        "movie_ids = movies.batch(1_000)\n",
        "unique_movie_ids = np.unique(np.concatenate(list(movie_ids)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vxFmdiVxT8j_"
      },
      "source": [
        "## Implementing a sequential model\n",
        "\n",
        "In our [basic retrieval tutorial](https://www.tensorflow.org/recommenders/examples/basic_retrieval), we use one query tower for the user, and the candidate tow for the candidate movie. However, the two-tower architecture is generalizble and not limited to \u003cuser,item\u003e pair. You can also use it to do item-to-item recommendation as we note in the [basic retrieval tutorial](https://www.tensorflow.org/recommenders/examples/basic_retrieval#item-to-item_recommendation).\n",
        "\n",
        "Here we are still going to use the two-tower architecture. Specificially, we use the query tower with a [Gated Recurrent Unit (GRU) layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/GRU) to encode the sequence of historical movies, and keep the same candidate tower for the candidate movie. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aP3OS_HPeefT"
      },
      "outputs": [],
      "source": [
        "embedding_dimension = 32\n",
        "\n",
        "query_model = tf.keras.Sequential([\n",
        "    tf.keras.layers.StringLookup(\n",
        "      vocabulary=unique_movie_ids, mask_token=None),\n",
        "    tf.keras.layers.Embedding(len(unique_movie_ids) + 1, embedding_dimension), \n",
        "    tf.keras.layers.GRU(embedding_dimension),\n",
        "])\n",
        "\n",
        "candidate_model = tf.keras.Sequential([\n",
        "  tf.keras.layers.StringLookup(\n",
        "      vocabulary=unique_movie_ids, mask_token=None),\n",
        "  tf.keras.layers.Embedding(len(unique_movie_ids) + 1, embedding_dimension)\n",
        "])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ich7YzXXAWq"
      },
      "source": [
        "The metrics, task and full model are defined similar to the basic retrieval model. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-KDYuhG9u9cq"
      },
      "outputs": [],
      "source": [
        "metrics = tfrs.metrics.FactorizedTopK(\n",
        "  candidates=movies.batch(128).map(candidate_model)\n",
        ")\n",
        "\n",
        "task = tfrs.tasks.Retrieval(\n",
        "  metrics=metrics\n",
        ")\n",
        "\n",
        "class Model(tfrs.Model):\n",
        "\n",
        "    def __init__(self, query_model, candidate_model):\n",
        "        super().__init__()\n",
        "        self._query_model = query_model\n",
        "        self._candidate_model = candidate_model\n",
        "\n",
        "        self._task = task\n",
        "\n",
        "    def compute_loss(self, features, training=False):\n",
        "        watch_history = features[\"context_movie_id\"]\n",
        "        watch_next_label = features[\"label_movie_id\"]\n",
        "\n",
        "        query_embedding = self._query_model(watch_history)       \n",
        "        candidate_embedding = self._candidate_model(watch_next_label)\n",
        "        \n",
        "        return self._task(query_embedding, candidate_embedding, compute_metrics=not training)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j22_GTTeXa_T"
      },
      "source": [
        "## Fitting and evaluating\n",
        "\n",
        "We can now compile, train and evaluate our sequential retrieval model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_c2wAqPVXYgz"
      },
      "outputs": [],
      "source": [
        "model = Model(query_model, candidate_model)\n",
        "model.compile(optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.1))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9Wto55LJWuj0"
      },
      "outputs": [],
      "source": [
        "cached_train = train_ds.shuffle(10_000).batch(12800).cache()\n",
        "cached_test = test_ds.batch(2560).cache()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DcXH5aDcWyJJ"
      },
      "outputs": [],
      "source": [
        "model.fit(cached_train, epochs=3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6aeXfXYw4Xtk"
      },
      "outputs": [],
      "source": [
        "model.evaluate(cached_test, return_dict=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ny2gMW3MdMKW"
      },
      "source": [
        "This concludes the sequential retrieval tutorial."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "sequential_retrieval.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
